#include <arpa/inet.h>
#include <errno.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/epoll.h>
#include <time.h>
#include <unistd.h>

#define MAX_EVENTS 1024
#define BUFFER_SIZE 1024
// #define PORT 4000

// 单核linux的极限pps（包每秒）一般在150k~300k, 再往上优化单核很难:  https://linux.cn/article-5724-1.html

void set_nonblocking(int sockfd)
{
    int opts = fcntl(sockfd, F_GETFL);
    if (opts < 0) {
        perror("fcntl(F_GETFL)");
        exit(EXIT_FAILURE);
    }
    opts = opts | O_NONBLOCK;
    if (fcntl(sockfd, F_SETFL, opts) < 0) {
        perror("fcntl(F_SETFL)");
        exit(EXIT_FAILURE);
    }
}

class MyTimer {
public:
    struct timespec start;
    struct timespec end;

    MyTimer()
    {
        clock_gettime(CLOCK_MONOTONIC, &start); // 开始计时
    }

    double elapsed()
    {
        clock_gettime(CLOCK_MONOTONIC, &end); // 开始计时
        double time_taken = (end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec) / 1e9;
        return time_taken;
    }
};

class EmuFpga {
public:
    void start(int port)
    {
        int udp_fd, epoll_fd;
        struct sockaddr_in serv_addr, cli_addr;
        struct epoll_event ev, events[MAX_EVENTS];
        char buffer[BUFFER_SIZE];
        socklen_t addr_len = sizeof(cli_addr);

        // 创建 UDP socket
        if ((udp_fd = socket(AF_INET, SOCK_DGRAM, 0)) < 0) {
            perror("socket");
            exit(EXIT_FAILURE);
        }

        set_nonblocking(udp_fd);

        // 设置服务器地址结构
        memset(&serv_addr, 0, sizeof(serv_addr));
        serv_addr.sin_family = AF_INET;
        serv_addr.sin_addr.s_addr = htonl(INADDR_ANY);
        serv_addr.sin_port = htons(port);

        // 绑定 socket
        if (bind(udp_fd, (struct sockaddr*)&serv_addr, sizeof(serv_addr)) < 0) {
            perror("bind");
            exit(EXIT_FAILURE);
        }

        // 创建 epoll 实例
        if ((epoll_fd = epoll_create1(0)) < 0) {
            perror("epoll_create1");
            exit(EXIT_FAILURE);
        }

        // 将 UDP socket 加入 epoll 监听列表
        ev.events = EPOLLIN;
        ev.data.fd = udp_fd;
        if (epoll_ctl(epoll_fd, EPOLL_CTL_ADD, udp_fd, &ev) < 0) {
            perror("epoll_ctl: udp_fd");
            exit(EXIT_FAILURE);
        }

        while (1) {
            int nfds = epoll_wait(epoll_fd, events, MAX_EVENTS, -1);
            if (nfds < 0) {
                perror("epoll_wait");
                exit(EXIT_FAILURE);
            }

            for (int i = 0; i < nfds; i++) {
                if (events[i].events & EPOLLIN) {
                    // 接收数据
                    int sock_fd = events[i].data.fd;
                    ssize_t n = recvfrom(sock_fd, buffer, BUFFER_SIZE, 0, (struct sockaddr*)&cli_addr, &addr_len);
                    if (n > 0) {
                        // 回显数据
                        sendto(sock_fd, buffer, n, 0, (struct sockaddr*)&cli_addr, addr_len);
                    } else if (n < 0 && errno != EAGAIN && errno != EWOULDBLOCK) {
                        perror("recvfrom");
                    }
                }
            }
        }

        close(udp_fd);
        close(epoll_fd);
    }
};

class PcSender {

    int udp_fd;
    int epoll_fd;
    struct sockaddr_in local_addr, fpga_addr;
    struct epoll_event ev, events[MAX_EVENTS];
    char send_buffer[BUFFER_SIZE];
    char recv_buffer[BUFFER_SIZE];
    socklen_t addr_len = sizeof(fpga_addr);

    int fpga_ip; // 远程数据的ip
    const char* fpga_ip_str;
    int fpga_port; // 远程数据的端口

    int local_ip; // 接收数据的ip
    const char* local_ip_str;
    int local_port; // 接收数据的端口
    MyTimer timer;

public:
    PcSender(const char* _fpga_ip_str, int _fpga_port)
    {
        fpga_ip_str = _fpga_ip_str;
        fpga_port = _fpga_port;

        local_ip_str = "127.0.0.1";
        local_port = fpga_port + 1;

        // 创建 UDP socket
        if ((udp_fd = socket(AF_INET, SOCK_DGRAM, 0)) < 0) {
            perror("socket");
            exit(EXIT_FAILURE);
        }

        set_nonblocking(udp_fd);

        // 设置服务器地址结构
        memset(&local_addr, 0, sizeof(local_addr));
        local_addr.sin_family = AF_INET;
        // 转换 IPv4 地址
        if (inet_pton(AF_INET, local_ip_str, &local_addr.sin_addr.s_addr) <= 0) {
            perror("inet_pton for IPv4");
            exit(EXIT_FAILURE);
        }
        // local_addr.sin_addr.s_addr = htonl(INADDR_ANY);
        local_addr.sin_port = htons(local_port);

        memset(&fpga_addr, 0, sizeof(fpga_addr));
        fpga_addr.sin_family = AF_INET;
        if (inet_pton(AF_INET, fpga_ip_str, &fpga_addr.sin_addr.s_addr) <= 0) {
            perror("inet_pton for IPv4");
            exit(EXIT_FAILURE);
        }
        // fpga_addr.sin_addr.s_addr = htonl(fpga_ip);
        fpga_addr.sin_port = htons(fpga_port);
    }

    void start()
    {
        // 绑定 socket
        if (bind(udp_fd, (struct sockaddr*)&local_addr, sizeof(local_addr)) < 0) {
            perror("bind");
            exit(EXIT_FAILURE);
        }

        // 创建 epoll 实例
        if ((epoll_fd = epoll_create1(0)) < 0) {
            perror("epoll_create1");
            exit(EXIT_FAILURE);
        }

        // 将 UDP socket 加入 epoll 监听列表
        ev.events = EPOLLIN;
        ev.data.fd = udp_fd;
        if (epoll_ctl(epoll_fd, EPOLL_CTL_ADD, udp_fd, &ev) < 0) {
            perror("epoll_ctl: udp_fd");
            exit(EXIT_FAILURE);
        }

        while (1) {
            do_send(udp_fd);
            int nfds = epoll_wait(epoll_fd, events, MAX_EVENTS, -1);
            if (nfds < 0) {
                perror("epoll_wait");
                exit(EXIT_FAILURE);
            }

            for (int i = 0; i < nfds; i++) {
                int sock_fd = events[i].data.fd;
                // if (events[i].events & EPOLLOUT) {
                //     if (sock_fd == udp_fd) {
                //         printf("EPOLLOUT.....\n");
                //         do_send(udp_fd);
                //     }
                // }
                if (events[i].events & EPOLLIN) {
                    do {
                        do_recive(sock_fd);
                    } while ((send_pkt - recv_pkt) > 100);
                }
            }
        }
        close(udp_fd);
        close(epoll_fd);
    }

    int send_pkt, recv_pkt, err_recv_pkt;
    long total_send_bytes, total_recv_bytes;

    void do_send(int sock_fd)
    {
        for (int i = 0; i < 200 - (send_pkt - recv_pkt); i++) {
            total_send_bytes += BUFFER_SIZE;
            send_pkt++;
            sendto(sock_fd, send_buffer, BUFFER_SIZE, 0, (struct sockaddr*)&fpga_addr, addr_len);
        }
    }

    struct sockaddr_in recv_fpga_addr; // 不用
    void do_recive(int sock_fd)
    {
        // 接收数据
        ssize_t n = recvfrom(sock_fd, recv_buffer, BUFFER_SIZE, 0, (struct sockaddr*)&recv_fpga_addr, &addr_len);
        if (n > 0) {
            recv_pkt++;
            total_recv_bytes += n;
            if (recv_pkt % 10000 == 0) {
                double elapsed_time = timer.elapsed();
                printf("recv...: send_pkt: %d, recv_pkt: %d, send_size: %.2fMB, recv_size: %.2fMB, speed: %.2fMB/s <-> s:%.2fKpps r:%.2fKpps,  rate: %.2f%, time: %.2fs\n",
                    send_pkt, recv_pkt,
                    float(total_send_bytes) / 1024.0 / 1024,
                    float(total_recv_bytes) / 1024.0 / 1024,
                    total_send_bytes / 1024.0 / 1024 / elapsed_time,
                    send_pkt / 1000 / elapsed_time,
                    recv_pkt / 1000 / elapsed_time,
                    float(recv_pkt) / send_pkt * 100,
                    elapsed_time);
                printf("reacv_size; %d\n", n);
            }
            // 回显数据
            // sendto(sock_fd, buffer, n, 0, (struct sockaddr*)&fpga_addr, addr_len);
        } else if (n < 0 && errno != EAGAIN && errno != EWOULDBLOCK) {
            err_recv_pkt++;
            perror("recvfrom");
        }
    }
};

int main(int argc, char* argv[])
{
    if (argc != 4) {
        printf("Usage: test [fpga|pc] ip port\n");
        return 1;
    }
    if (strcmp("fpga", argv[1]) == 0) {
        EmuFpga().start(std::atoi(argv[3]));
    } else if (strcmp("pc", argv[1]) == 0) {
        PcSender(argv[2], std::atoi(argv[3])).start();
    } else {
        printf("Error: command error [%s]\n", argv[1]);
        return 1;
    }
    return 0;
}
